# 知识图谱服务
import networkx as nx
from collections import defaultdict


class KnowledgeGraph:
    def __init__(self, dataset):
        self.graph = nx.DiGraph()
        self._build_graph(dataset.data)

    def _build_graph(self, data):
        """从论文数据构建知识图谱"""
        category_counts = defaultdict(int)

        # 第一层：学科领域节点
        for category in data['categories']:
            main_field = category.split()[0]
            category_counts[main_field] += 1

        # 添加学科领域节点
        for field, count in category_counts.items():
            self.graph.add_node(field, type="field", count=count)

        # 添加论文节点和关系
        for idx, row in data.iterrows():
            paper_id = f"paper_{idx}"
            main_field = row['categories'].split()[0]

            self.graph.add_node(paper_id,
                                type="paper",
                                title=row['title'],
                                abstract=row['abstract']
                                )
            self.graph.add_edge(paper_id, main_field, relation="belongs_to")
            self.graph.add_edge(main_field, paper_id, relation="contains")

    def get_related_papers(self, field, top_n=10):
        """获取领域相关论文"""
        return [n for n in self.graph.neighbors(field)
                if self.graph.nodes[n]['type'] == 'paper'][:top_n]